from typing import Any, Dict, List
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain
from langchain_community.llms import DeepSeek
from ..base import BasePredictionModule


class IntentAnalyzer:
    """意图分析器"""
    def __init__(self, llm):
        self.llm = llm
        self.template = """分析以下查询的意图：

查询: {query}

请提供以下信息：
1. 主要意图
2. 查询类型（如：事实查询、分析查询、比较查询等）
3. 预期答案类型
4. 关键实体
5. 时间范围（如果有）
6. 其他相关上下文

输出格式：
{
    "intent": "主要意图",
    "query_type": "查询类型",
    "answer_type": "预期答案类型",
    "entities": ["实体1", "实体2"],
    "time_range": "时间范围",
    "context": "相关上下文"
}"""
        
        self.prompt = PromptTemplate(
            template=self.template,
            input_variables=["query"]
        )
        self.chain = LLMChain(llm=llm, prompt=self.prompt)

    async def analyze(self, query: str) -> Dict[str, Any]:
        """分析查询意图"""
        result = await self.chain.arun(query=query)
        return eval(result)  # 注意：在生产环境中应使用更安全的JSON解析


class ContextGenerator:
    """上下文生成器"""
    def __init__(self, llm):
        self.llm = llm
        self.template = """基于以下查询意图，生成可能的上下文信息：

查询: {query}
意图分析: {intent}

请生成可能的相关上下文，包括：
1. 背景信息
2. 相关概念
3. 潜在的关联信息
4. 可能的限制条件

生成的上下文："""
        
        self.prompt = PromptTemplate(
            template=self.template,
            input_variables=["query", "intent"]
        )
        self.chain = LLMChain(llm=llm, prompt=self.prompt)

    async def generate(self, query: str, intent: Dict[str, Any]) -> str:
        """生成上下文"""
        return await self.chain.arun(query=query, intent=str(intent))


class QueryRewriter:
    """查询重写器"""
    def __init__(self, llm):
        self.llm = llm
        self.template = """基于以下查询和上下文，生成多个改写的查询：

原始查询: {query}
上下文: {context}

请生成3个不同的查询表达，确保：
1. 保持原始意图
2. 使用不同的表达方式
3. 添加上下文中的关键信息
4. 消除歧义

改写的查询："""
        
        self.prompt = PromptTemplate(
            template=self.template,
            input_variables=["query", "context"]
        )
        self.chain = LLMChain(llm=llm, prompt=self.prompt)

    async def rewrite(self, query: str, context: str) -> List[str]:
        """重写查询"""
        result = await self.chain.arun(query=query, context=context)
        return [q.strip() for q in result.split("\n") if q.strip()]


class PredictionModule(BasePredictionModule):
    """预测模块实现"""
    def __init__(self):
        self.llm = None
        self.intent_analyzer = None
        self.context_generator = None
        self.query_rewriter = None
        self.config = None

    def initialize(self, config: Dict[str, Any]) -> None:
        """初始化模块"""
        self.config = config
        self.llm = DeepSeek(
            model=config["model"]["name"],
            temperature=config["model"]["temperature"],
            max_tokens=config["model"]["max_tokens"]
        )
        
        self.intent_analyzer = IntentAnalyzer(self.llm)
        self.context_generator = ContextGenerator(self.llm)
        self.query_rewriter = QueryRewriter(self.llm)

    def validate_config(self, config: Dict[str, Any]) -> bool:
        """验证配置"""
        required = ["model", "intent_threshold", "context_window"]
        if not all(key in config for key in required):
            return False
        
        model_required = ["name", "temperature", "max_tokens"]
        return all(key in config["model"] for key in model_required)

    async def analyze_intent(self, query: str) -> Dict[str, Any]:
        """分析查询意图"""
        if not self.intent_analyzer:
            raise RuntimeError("Module not initialized")
        
        intent = await self.intent_analyzer.analyze(query)
        
        # 验证意图置信度
        if intent.get("confidence", 1.0) < self.config["intent_threshold"]:
            intent["needs_clarification"] = True
            
        return intent

    async def generate_context(self, query: str, intent: Dict[str, Any]) -> str:
        """生成上下文"""
        if not self.context_generator:
            raise RuntimeError("Module not initialized")
            
        context = await self.context_generator.generate(query, intent)
        
        # 截断上下文以适应窗口大小
        if len(context) > self.config["context_window"]:
            context = context[:self.config["context_window"]]
            
        return context

    async def rewrite_query(self, query: str, context: str) -> List[str]:
        """重写查询"""
        if not self.query_rewriter:
            raise RuntimeError("Module not initialized")
            
        if not self.config["query_rewrite"]["enabled"]:
            return [query]
            
        rewrites = await self.query_rewriter.rewrite(query, context)
        
        # 限制重写数量
        max_rewrites = self.config["query_rewrite"]["max_rewrites"]
        return rewrites[:max_rewrites] 